from langchain.docstore.document import Document
from configs import EMBEDDING_MODEL, logger
from server.model_workers.base import ApiEmbeddingsParams
from server.utils import BaseResponse, get_model_worker_config, list_embed_models, list_online_embed_models
from fastapi import Body
from fastapi.concurrency import run_in_threadpool
from typing import Dict, List
import requests

online_embed_models = list_online_embed_models()

'''
embedding txt use post
curl -X 'POST' \
  'http://127.0.0.1:8501/v1/embeddings' \
  -H 'accept: application/json' \
  -H 'Content-Type: application/json' \
  -d '{
  "input": "to be embedding str"
}'


schema:
"/v1/embeddings": {
            "post": {
                "summary": "Embedding",
                "operationId": "create_embedding_v1_embeddings_post",
                "requestBody": {
                    "content": {
                        "application/json": {
                            "schema": {
                                "$ref": "#/components/schemas/CreateEmbeddingRequest"
                            }
                        }
                    },
                    "required": true
                },
                "responses": {
                    "200": {
                        "description": "Successful Response",
                        "content": {
                            "application/json": {
                                "schema": {}
                            }
                        }
                    },
                    "422": {
                        "description": "Validation Error",
                        "content": {
                            "application/json": {
                                "schema": {
                                    "$ref": "#/components/schemas/HTTPValidationError"
                                }
                            }
                        }
                    }
                },
                "security": [
                    {
                        "HTTPBearer": []
                    }
                ]
            }
        },
        
 "CreateEmbeddingRequest": {
                "properties": {
                    "model": {
                        "anyOf": [
                            {
                                "type": "string"
                            },
                            {
                                "type": "null"
                            }
                        ],
                        "title": "Model",
                        "description": "The model to use for generating completions."
                    },
                    "input": {
                        "anyOf": [
                            {
                                "type": "string"
                            },
                            {
                                "items": {
                                    "type": "string"
                                },
                                "type": "array"
                            }
                        ],
                        "title": "Input",
                        "description": "The input to embed."
                    },
                    "user": {
                        "anyOf": [
                            {
                                "type": "string"
                            },
                            {
                                "type": "null"
                            }
                        ],
                        "title": "User"
                    }
                },
                "type": "object",
                "required": [
                    "input"
                ],
                "title": "CreateEmbeddingRequest",
                "examples": [
                    {
                        "input": "The food was delicious and the waiter..."
                    }
                ]
            },
            "HTTPValidationError": {
                "properties": {
                    "detail": {
                        "items": {
                            "$ref": "#/components/schemas/ValidationError"
                        },
                        "type": "array",
                        "title": "Detail"
                    }
                },
                "type": "object",
                "title": "HTTPValidationError"
            },
'''


def openai_embedding(texts: List[str]) -> BaseResponse(data=List[List[float]]):
    """Embeds a list of texts using the specified endpoint and API key.

    Args:
        texts (List[str]): A list of texts to embed.

    Returns:
        List[dict]: A list of embedding dictionaries, one for each text. Each dictionary
                    contains the embedded vector ("embedding") and any additional response
                    data (e.g., "completion", "model") according to the provided schema.

    Raises:
        ValueError: If the response status code is not 200.
        RuntimeError: If any other error occurs.
    """

    api_url = "http://127.0.0.1:8501/v1/embeddings"
    api_key = "NULL"  # Replace with your API key

    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }

    body = {
        "input": texts,
    }

    try:
        response = requests.post(api_url, headers=headers, json=body)
        response.raise_for_status()  # Raise an exception for non-200 status codes

        response_data = response.json()

        if isinstance(response_data, list):
            # Assume list of embeddings if schema is not provided
            embeddings = response_data
        else:
            # Handle dictionary response based on provided schema
            embeddings = [ item["embedding"] for item in response_data['data'] ]

        return BaseResponse(data=embeddings)

    except (requests.exceptions.RequestException, ValueError) as e:
        raise RuntimeError(f"Error embedding texts: {e}")


# Example usage
'''
texts = ["to be embedding str", "another text to embed"]
embeddings = openai_embedding(texts)
print(embeddings)
'''

def embed_texts(
        texts: List[str],
        embed_model: str = EMBEDDING_MODEL,
        to_query: bool = False,
) -> BaseResponse:
    '''
    对文本进行向量化。返回数据格式：BaseResponse(data=List[List[float]])
    '''


    try:
        if embed_model.find('text-embedding-ada-002') != -1:
            return openai_embedding(texts)

        if embed_model in list_embed_models():  # 使用本地Embeddings模型
            from server.utils import load_local_embeddings

            embeddings = load_local_embeddings(model=embed_model)
            return BaseResponse(data=embeddings.embed_documents(texts))

        if embed_model in list_online_embed_models():  # 使用在线API
            config = get_model_worker_config(embed_model)
            worker_class = config.get("worker_class")
            embed_model = config.get("embed_model")
            worker = worker_class()
            if worker_class.can_embedding():
                params = ApiEmbeddingsParams(texts=texts, to_query=to_query, embed_model=embed_model)
                resp = worker.do_embeddings(params)
                return BaseResponse(**resp)

        return BaseResponse(code=500, msg=f"指定的模型 {embed_model} 不支持 Embeddings 功能。")
    except Exception as e:
        logger.error(e)
        return BaseResponse(code=500, msg=f"文本向量化过程中出现错误：{e}")


async def aembed_texts(
    texts: List[str],
    embed_model: str = EMBEDDING_MODEL,
    to_query: bool = False,
) -> BaseResponse:
    '''
    对文本进行向量化。返回数据格式：BaseResponse(data=List[List[float]])
    '''
    try:
        if embed_model in list_embed_models(): # 使用本地Embeddings模型
            from server.utils import load_local_embeddings

            embeddings = load_local_embeddings(model=embed_model)
            return BaseResponse(data=await embeddings.aembed_documents(texts))

        if embed_model in list_online_embed_models(): # 使用在线API
            return await run_in_threadpool(embed_texts,
                                           texts=texts,
                                           embed_model=embed_model,
                                           to_query=to_query)
    except Exception as e:
        logger.error(e)
        return BaseResponse(code=500, msg=f"文本向量化过程中出现错误：{e}")


def embed_texts_endpoint(
        texts: List[str] = Body(..., description="要嵌入的文本列表", examples=[["hello", "world"]]),
        embed_model: str = Body(EMBEDDING_MODEL,
                                description=f"使用的嵌入模型，除了本地部署的Embedding模型，也支持在线API({online_embed_models})提供的嵌入服务。"),
        to_query: bool = Body(False, description="向量是否用于查询。有些模型如Minimax对存储/查询的向量进行了区分优化。"),
) -> BaseResponse:
    '''
    对文本进行向量化，返回 BaseResponse(data=List[List[float]])
    '''
    return embed_texts(texts=texts, embed_model=embed_model, to_query=to_query)


def embed_documents(
        docs: List[Document],
        embed_model: str = EMBEDDING_MODEL,
        to_query: bool = False,
) -> Dict:
    """
    将 List[Document] 向量化，转化为 VectorStore.add_embeddings 可以接受的参数
    """
    texts = [x.page_content for x in docs]
    metadatas = [x.metadata for x in docs]
    embeddings = embed_texts(texts=texts, embed_model=embed_model, to_query=to_query).data
    if embeddings is not None:
        return {
            "texts": texts,
            "embeddings": embeddings,
            "metadatas": metadatas,
        }
